batcher_improvements: added flushing and some extra checks

Author: brianshannan-wf

batcher/batcher.go

@@ -1,3 +1,19 @@
+/*
+Copyright 2015 Workiva, LLC
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
 package batcher
 
 import (
@@ -15,9 +31,15 @@ type Batcher interface {
 	// one of the conditions for a "complete" batch is reached.
 	Get() ([]interface{}, error)
 
+	// Flush forcibly completes the batch currently being built
+	Flush() error
+
 	// Dispose will dispose of the batcher. Any calls to Put or Get
 	// will return errors.
 	Dispose()
+
+	// IsDisposed will determine if the batcher is disposed
+	IsDisposed() bool
 }
 
 // ErrDisposed is the error returned for a disposed Batcher
@@ -30,7 +52,6 @@ type basicBatcher struct {
 	maxTime        time.Duration
 	maxItems       uint
 	maxBytes       uint
-	queueLen       uint
 	calculateBytes CalculateBytes
 	disposed       bool
 	items          []interface{}
@@ -55,7 +76,6 @@ func New(maxTime time.Duration, maxItems, maxBytes, queueLen uint, calculate Cal
 		maxTime:        maxTime,
 		maxItems:       maxItems,
 		maxBytes:       maxBytes,
-		queueLen:       queueLen,
 		calculateBytes: calculate,
 		items:          make([]interface{}, 0, maxItems),
 		batchChan:      make(chan []interface{}, queueLen),
@@ -71,11 +91,11 @@ func (b *basicBatcher) Put(item interface{}) error {
 	}
 
 	b.items = append(b.items, item)
-	b.availableBytes += b.calculateBytes(item)
+	if b.calculateBytes != nil {
+		b.availableBytes += b.calculateBytes(item)
+	}
 	if b.ready() {
-		b.batchChan <- b.items
-		b.items = make([]interface{}, 0, b.maxItems)
-		b.availableBytes = 0
+		b.flush()
 	}
 
 	b.lock.Unlock()
@@ -85,12 +105,8 @@ 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.
 func (b *basicBatcher) Get() ([]interface{}, error) {
-	b.lock.RLock()
-	if b.disposed {
-		b.lock.RUnlock()
-		return nil, ErrDisposed
-	}
-	b.lock.RUnlock()
+	// Don't check disposed yet so any items remaining in the queue
+	// will be returned properly.
 
 	var timeout <-chan time.Time
 	if b.maxTime > 0 {
@@ -117,16 +133,49 @@ 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()
+		return ErrDisposed
+	}
+	b.flush()
+	b.lock.Unlock()
+	return nil
+}
+
 // Dispose will dispose of the batcher. Any calls to Put or Get
 // will return errors.
 func (b *basicBatcher) Dispose() {
 	b.lock.Lock()
+	if b.disposed {
+		b.lock.Unlock()
+		return
+	}
+	b.flush()
 	b.disposed = true
 	b.items = nil
 	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
+}
+
+// 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
+	b.items = make([]interface{}, 0, b.maxItems)
+	b.availableBytes = 0
+}
+
 func (b *basicBatcher) ready() bool {
 	if b.maxItems != 0 && uint(len(b.items)) >= b.maxItems {
 		return true

batcher/batcher_test.go

@@ -1,3 +1,19 @@
+/*
+Copyright 2015 Workiva, LLC
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
 package batcher
 
 import (
@@ -72,6 +88,25 @@ func TestMaxTime(t *testing.T) {
 	assert.Nil(err)
 }
 
+func TestFlush(t *testing.T) {
+	assert := assert.New(t)
+	b, err := New(0, 10, 10, 10, func(str interface{}) uint {
+		return uint(len(str.(string)))
+	})
+	assert.Nil(err)
+	b.Put("a")
+	wait := make(chan bool)
+	go func() {
+		batch, err := b.Get()
+		assert.Equal([]interface{}{"a"}, batch)
+		assert.Nil(err)
+		wait <- true
+	}()
+
+	b.Flush()
+	<-wait
+}
+
 func TestMultiConsumer(t *testing.T) {
 	assert := assert.New(t)
 	b, err := New(0, 100, 100000, 10, func(str interface{}) uint {
@@ -101,21 +136,41 @@ func TestMultiConsumer(t *testing.T) {
 
 func TestDispose(t *testing.T) {
 	assert := assert.New(t)
-	b, err := New(0, 100000, 100000, 10, func(str interface{}) uint {
+	b, err := New(0, 2, 100000, 10, 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() {
-		_, err := b.Get()
+		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
 	}()
 
 	b.Dispose()
 
-	assert.Equal(ErrDisposed, b.Put("a"))
+	assert.Equal(ErrDisposed, b.Put("d"))
+	assert.Equal(ErrDisposed, b.Flush())
 
 	<-wait
 }
+
+func TestIsDisposed(t *testing.T) {
+	assert := assert.New(t)
+	b, err := New(0, 10, 10, 10, func(str interface{}) uint {
+		return uint(len(str.(string)))
+	})
+	assert.Nil(err)
+	assert.False(b.IsDisposed())
+	b.Dispose()
+	assert.True(b.IsDisposed())
+}

mock/batcher.go

@@ -1,9 +1,29 @@
+/*
+Copyright 2015 Workiva, LLC
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
 package mock
 
 import (
 	"github.com/stretchr/testify/mock"
+
+	"github.com/Workiva/go-datastructures/batcher"
 )
 
+var _ batcher.Batcher = new(Batcher)
+
 type Batcher struct {
 	mock.Mock
 	PutChan chan bool
@@ -22,6 +42,16 @@ func (m *Batcher) Get() ([]interface{}, error) {
 	return args.Get(0).([]interface{}), args.Error(1)
 }
 
+func (m *Batcher) Flush() error {
+	args := m.Called()
+	return args.Error(0)
+}
+
 func (m *Batcher) Dispose() {
 	m.Called()
 }
+
+func (m *Batcher) IsDisposed() bool {
+	args := m.Called()
+	return args.Bool(0)
+}