pubsub: Ensure batch flushes on shutdown even if MinBatchSize isn't met (#3543)

Author: vangent

pubsub/batcher/batcher.go

@@ -91,6 +91,7 @@ type Options struct {
 	// Maximum number of concurrent handlers. Defaults to 1.
 	MaxHandlers int
 	// Minimum size of a batch. Defaults to 1.
+	// May be ignored during shutdown.
 	MinBatchSize int
 	// Maximum size of a batch. 0 means no limit.
 	MaxBatchSize int
@@ -199,26 +200,33 @@ func (b *Batcher) AddNoWait(item any) <-chan error {
 	b.pending = append(b.pending, waiter{item, c})
 	if b.nHandlers < 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()
-		if batch != nil {
-			b.wg.Add(1)
-			go func() {
-				b.callHandler(batch)
-				b.wg.Done()
-			}()
-			b.nHandlers++
-		}
+		b.handleBatch(b.nextBatch())
 	}
 	// If we can't start a handler, then one of the currently running handlers will
 	// take our item.
 	return c
 }
 
+// Requires b.mu be held.
+func (b *Batcher) handleBatch(batch []waiter) {
+	if len(batch) == 0 {
+		return
+	}
+	b.wg.Add(1)
+	go func() {
+		b.callHandler(batch)
+		b.wg.Done()
+	}()
+	b.nHandlers++
+}
+
 // nextBatch returns the batch to process, and updates b.pending.
 // It returns nil if there's no batch ready for processing.
 // b.mu must be held.
 func (b *Batcher) nextBatch() []waiter {
-	if len(b.pending) < b.opts.MinBatchSize {
+	// If we're not shutting down, respect minimums.  If we're shutting down
+	// though, we ignore minimums to make sure everything is flushed.
+	if !b.shutdown && len(b.pending) < b.opts.MinBatchSize {
 		return nil
 	}
 
@@ -282,6 +290,12 @@ func (b *Batcher) callHandler(batch []waiter) {
 func (b *Batcher) Shutdown() {
 	b.mu.Lock()
 	b.shutdown = true
+	// If there aren't any handlers running, there might be a partial
+	// batch. Make sure it gets flushed even if it hasn't reached the
+	// minimums.
+	if b.nHandlers == 0 {
+		b.handleBatch(b.nextBatch())
+	}
 	b.mu.Unlock()
 	b.wg.Wait()
 }

pubsub/batcher/batcher_test.go

@@ -271,6 +271,33 @@ func TestShutdown(t *testing.T) {
 	}
 }
 
+// TestMinBatchSizeFlushesOnShutdown ensures that Shutdown() flushes batches, even if
+// the pending count is less than the minimum batch size.
+func TestMinBatchSizeFlushesOnShutdown(t *testing.T) {
+	var got [][]int
+
+	batchSize := 3
+	b := batcher.New(reflect.TypeOf(int(0)), &batcher.Options{MinBatchSize: batchSize}, func(items interface{}) error {
+		got = append(got, items.([]int))
+		return nil
+	})
+	for i := 0; i < (batchSize - 1); i++ {
+		b.AddNoWait(i)
+	}
+
+	// Ensure that we've received nothing
+	if len(got) > 0 {
+		t.Errorf("got batch unexpectedly: %+v", got)
+	}
+
+	b.Shutdown()
+
+	want := [][]int{{0, 1}}
+	if !cmp.Equal(got, want) {
+		t.Errorf("got %+v, want %+v on shutdown", got, want)
+	}
+}
+
 func TestItemCanBeInterface(t *testing.T) {
 	readerType := reflect.TypeOf([]io.Reader{}).Elem()
 	called := false