Atomically update nHandlers

pubsub/batcher/batcher.go

@@ -23,6 +23,7 @@ import (
 	"errors"
 	"reflect"
 	"sync"
+	"sync/atomic"
 )
 
 // Split determines how to split n (representing n items) into batches based on
@@ -70,7 +71,7 @@ type Batcher struct {
 
 	mu        sync.Mutex
 	pending   []waiter // items waiting to be handled
-	nHandlers int      // number of currently running handler goroutines
+	nHandlers int32    // number of currently running handler goroutines
 	shutdown  bool
 }
 
@@ -197,7 +198,7 @@ func (b *Batcher) AddNoWait(item interface{}) <-chan error {
 
 	// Add the item to the pending list.
 	b.pending = append(b.pending, waiter{item, c})
-	if b.nHandlers < b.opts.MaxHandlers {
+	if atomic.LoadInt32(&b.nHandlers) < int32(b.opts.MaxHandlers) {
 		// If we can start a handler, do so with the item just added and any others that are pending.
 		batch := b.nextBatch()
 		b.handleBatch(batch)
@@ -217,7 +218,7 @@ func (b *Batcher) handleBatch(batch []waiter) {
 		b.callHandler(batch)
 		b.wg.Done()
 	}()
-	b.nHandlers++
+	atomic.AddInt32(&b.nHandlers, 1)
 }
 
 // nextBatch returns the batch to process, and updates b.pending.
@@ -284,7 +285,7 @@ func (b *Batcher) callHandler(batch []waiter) {
 		// always get to run.
 		batch = b.nextBatch()
 		if batch == nil {
-			b.nHandlers--
+			atomic.AddInt32(&b.nHandlers, -1)
 		}
 		b.mu.Unlock()
 	}
@@ -300,7 +301,7 @@ func (b *Batcher) Shutdown() {
 
 	// On shutdown, ensure that we attempt to flush any pending items
 	// if there's a minimum batch size.
-	if b.nHandlers < b.opts.MaxHandlers {
+	if atomic.LoadInt32(&b.nHandlers) < int32(b.opts.MaxHandlers) {
 		batch := b.nextBatch()
 		b.handleBatch(batch)
 	}