chipingress/batch: publish queued messages on Stop instead of discarding them

pkg/chipingress/batch/batcher.go

@@ -6,7 +6,8 @@ import (
 )
 
 // batchWithInterval reads from a channel and calls processFn with batches based on size or time interval.
-// When the context is cancelled, it automatically flushes any remaining items in the batch.
+// When the context is cancelled, it flushes the current batch and every item already queued on input,
+// in batches of at most batchSize, then returns.
 func batchWithInterval[T any](
 	ctx context.Context,
 	input <-chan T,
@@ -29,7 +30,8 @@ func batchWithInterval[T any](
 	for {
 		select {
 		case <-ctx.Done():
-			flush()
+			timer.Stop()
+			drain(input, batch, batchSize, processFn)
 			return
 		case msg, ok := <-input:
 			if !ok {
@@ -56,3 +58,18 @@ func batchWithInterval[T any](
 		}
 	}
 }
+
+// drain processes batch plus everything already buffered on input, without waiting for new items.
+// The batcher is input's only reader, so a non-zero len guarantees the receive won't block.
+func drain[T any](input <-chan T, batch []T, batchSize int, processFn func([]T)) {
+	for len(input) > 0 {
+		batch = append(batch, <-input)
+		if len(batch) >= batchSize {
+			processFn(batch)
+			batch = nil
+		}
+	}
+	if len(batch) > 0 {
+		processFn(batch)
+	}
+}

pkg/chipingress/batch/client.go

@@ -103,14 +103,18 @@ type Client struct {
 	maxPublishTimeout       time.Duration
 	messageBuffer           chan *messageWithCallback
 	stopCh                  stopCh
-	log                     *zap.SugaredLogger
-	callbackWg              sync.WaitGroup
-	shutdownTimeout         time.Duration
-	shutdownOnce            sync.Once
-	batcherDone             chan struct{}
-	started                 bool
-	counters                sync.Map // map[seqnumKey]*atomic.Uint64 for per-(source,type) seqnum, cleared on Stop()
-	clientName              string
+	// admission makes "check stopCh, then enqueue" atomic with respect to Stop closing stopCh,
+	// so every message QueueMessage accepted is on messageBuffer before the batcher drains it.
+	admission       sync.RWMutex
+	afterFastPath   func() // test-only hook between the fast-path stop check and admission; nil in production
+	log             *zap.SugaredLogger
+	callbackWg      sync.WaitGroup
+	shutdownTimeout time.Duration
+	shutdownOnce    sync.Once
+	batcherDone     chan struct{}
+	started         bool
+	counters        sync.Map // map[seqnumKey]*atomic.Uint64 for per-(source,type) seqnum, cleared on Stop()
+	clientName      string
 
 	metrics batchClientMetrics
 
@@ -205,18 +209,21 @@ func (b *Client) Start(ctx context.Context) {
 }
 
 // Stop ensures:
-// - current batch is flushed
+// - every message already queued is published, in batches of at most the batch size
 // - all current network calls are completed
 // - all callbacks are completed
-// Forcibly shutdowns down after timeout if not completed.
+// Forcibly shuts down after the shutdown timeout if not completed; batches still being
+// drained then fail against the closed client and report the error to their callbacks.
 func (b *Client) Stop() {
 	b.shutdownOnce.Do(func() {
 		// Use a standalone timeout context so the shutdown wait isn't cancelled
 		// by close(b.stopCh) below.
 		ctx, cancel := context.WithTimeout(context.Background(), b.shutdownTimeout)
 		defer cancel()
 
+		b.admission.Lock()
 		close(b.stopCh)
+		b.admission.Unlock()
 
 		// Only wait for the batcher goroutine when Start() was called;
 		// otherwise batcherDone is never closed and we'd block until timeout.
@@ -279,12 +286,15 @@ func (b *Client) QueueMessage(event *chipingress.CloudEventPb, callback func(err
 		return nil
 	}
 
-	// Check shutdown first to avoid race with buffer send
+	// Fast path: skip cloning and seqnum stamping once shut down. Rechecked under admission below.
 	select {
 	case <-b.stopCh:
 		return ErrClientShutdown
 	default:
 	}
+	if b.afterFastPath != nil {
+		b.afterFastPath()
+	}
 
 	eventToQueue := event
 	if b.cloneEvent {
@@ -312,6 +322,13 @@ func (b *Client) QueueMessage(event *chipingress.CloudEventPb, callback func(err
 		callback: callback,
 	}
 
+	b.admission.RLock()
+	defer b.admission.RUnlock()
+	select {
+	case <-b.stopCh:
+		return ErrClientShutdown
+	default:
+	}
 	select {
 	case b.messageBuffer <- msg:
 		return nil

pkg/chipingress/batch/client_test.go

@@ -7,6 +7,7 @@ import (
 	"strconv"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -18,6 +19,7 @@ import (
 	"go.opentelemetry.io/otel/attribute"
 	sdkmetric "go.opentelemetry.io/otel/sdk/metric"
 	"go.opentelemetry.io/otel/sdk/metric/metricdata"
+	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
 	"google.golang.org/protobuf/proto"
@@ -1215,6 +1217,185 @@ func TestStop(t *testing.T) {
 
 		assert.Equal(t, 0, countCounters(&client.counters))
 	})
+
+	t.Run("Stop publishes messages still queued behind in-flight sends", func(t *testing.T) {
+		const queued = 100
+		mockClient := mocks.NewClient(t)
+		mockClient.EXPECT().Close().Return(nil).Maybe()
+
+		firstSendStarted := make(chan struct{})
+		release := make(chan struct{})
+		var mu sync.Mutex
+		published := map[string]bool{}
+		mockClient.EXPECT().PublishBatch(mock.Anything, mock.Anything).
+			RunAndReturn(func(_ context.Context, batch *chipingress.CloudEventBatch, _ ...grpc.CallOption) (*chipingress.PublishResponse, error) {
+				mu.Lock()
+				first := len(published) == 0
+				results := make([]*chipingress.PublishResult, len(batch.Events))
+				for i, e := range batch.Events {
+					published[e.Id] = true
+					results[i] = &chipingress.PublishResult{EventId: e.Id}
+				}
+				mu.Unlock()
+				if first {
+					close(firstSendStarted)
+					<-release // hold the only send slot so the rest of the messages stay queued
+				}
+				return &chipingress.PublishResponse{Results: results}, nil
+			})
+
+		client, err := NewBatchClient(mockClient,
+			WithBatchSize(2),
+			WithMaxConcurrentSends(1),
+			WithMessageBuffer(queued),
+			WithBatchInterval(time.Hour),
+			WithShutdownTimeout(5*time.Second),
+		)
+		require.NoError(t, err)
+		client.Start(t.Context())
+
+		callbackCount := 0
+		var callbackErrs []error
+		for i := range queued {
+			require.NoError(t, client.QueueMessage(&chipingress.CloudEventPb{
+				Id: strconv.Itoa(i), Source: "test-source", Type: "test.event.type",
+			}, func(err error) {
+				mu.Lock()
+				defer mu.Unlock()
+				callbackCount++
+				if err != nil {
+					callbackErrs = append(callbackErrs, err)
+				}
+			}))
+		}
+		<-firstSendStarted
+
+		stopped := make(chan struct{})
+		go func() {
+			client.Stop()
+			close(stopped)
+		}()
+		require.Eventually(t, func() bool {
+			select {
+			case <-client.stopCh:
+				return true
+			default:
+				return false
+			}
+		}, time.Second, time.Millisecond)
+		close(release)
+		<-stopped // Stop waits for in-flight sends and callbacks
+
+		mu.Lock()
+		defer mu.Unlock()
+		assert.Len(t, published, queued, "every queued message must be published on Stop")
+		assert.Equal(t, queued, callbackCount, "every queued message must get its callback")
+		assert.Empty(t, callbackErrs)
+	})
+
+	t.Run("Stop returns at the shutdown timeout and late batches report errors", func(t *testing.T) {
+		const queued = 20
+		mockClient := mocks.NewClient(t)
+		var closed atomic.Bool
+		mockClient.EXPECT().Close().RunAndReturn(func() error { closed.Store(true); return nil }).Once()
+
+		firstSendStarted := make(chan struct{})
+		release := make(chan struct{})
+		var firstOnce sync.Once
+		mockClient.EXPECT().PublishBatch(mock.Anything, mock.Anything).
+			RunAndReturn(func(_ context.Context, batch *chipingress.CloudEventBatch, _ ...grpc.CallOption) (*chipingress.PublishResponse, error) {
+				if closed.Load() {
+					return nil, errors.New("client connection is closing")
+				}
+				firstOnce.Do(func() {
+					close(firstSendStarted)
+					<-release // outlive the shutdown timeout
+				})
+				results := make([]*chipingress.PublishResult, len(batch.Events))
+				for i, e := range batch.Events {
+					results[i] = &chipingress.PublishResult{EventId: e.Id}
+				}
+				return &chipingress.PublishResponse{Results: results}, nil
+			})
+
+		const shutdownTimeout = 100 * time.Millisecond
+		client, err := NewBatchClient(mockClient,
+			WithBatchSize(2),
+			WithMaxConcurrentSends(1),
+			WithMessageBuffer(queued),
+			WithBatchInterval(time.Hour),
+			WithShutdownTimeout(shutdownTimeout),
+		)
+		require.NoError(t, err)
+		client.Start(t.Context())
+
+		var mu sync.Mutex
+		var ok, failed int
+		for i := range queued {
+			require.NoError(t, client.QueueMessage(&chipingress.CloudEventPb{
+				Id: strconv.Itoa(i), Source: "test-source", Type: "test.event.type",
+			}, func(err error) {
+				mu.Lock()
+				defer mu.Unlock()
+				if err != nil {
+					failed++
+				} else {
+					ok++
+				}
+			}))
+		}
+		<-firstSendStarted
+
+		start := time.Now()
+		client.Stop()
+		elapsed := time.Since(start)
+		assert.GreaterOrEqual(t, elapsed, shutdownTimeout)
+		assert.Less(t, elapsed, shutdownTimeout+time.Second, "Stop must return at the shutdown timeout")
+		require.True(t, closed.Load(), "Stop must close the client after the timeout")
+
+		close(release)
+		require.EventuallyWithT(t, func(c *assert.CollectT) {
+			mu.Lock()
+			defer mu.Unlock()
+			assert.Equal(c, queued, ok+failed, "every queued message must get its callback")
+		}, 5*time.Second, 10*time.Millisecond)
+		mu.Lock()
+		defer mu.Unlock()
+		assert.Equal(t, 2, ok, "only the batch in flight before the timeout succeeds")
+		assert.Equal(t, queued-2, failed, "batches drained after the timeout report the closed-client error")
+	})
+
+	t.Run("QueueMessage that passed the fast-path check before Stop is rejected, not abandoned", func(t *testing.T) {
+		mockClient := mocks.NewClient(t)
+		mockClient.EXPECT().Close().Return(nil).Maybe()
+
+		client, err := NewBatchClient(mockClient, WithBatchInterval(time.Hour))
+		require.NoError(t, err)
+		client.Start(t.Context())
+
+		reachedHook := make(chan struct{})
+		resume := make(chan struct{})
+		client.afterFastPath = func() {
+			close(reachedHook)
+			<-resume // hold the producer between the fast-path check and admission
+		}
+
+		var callbackRan atomic.Bool
+		queueErr := make(chan error, 1)
+		go func() {
+			queueErr <- client.QueueMessage(&chipingress.CloudEventPb{
+				Id: "late", Source: "test-source", Type: "test.event.type",
+			}, func(error) { callbackRan.Store(true) })
+		}()
+		<-reachedHook
+
+		client.Stop() // drains an empty buffer and returns while the producer is parked
+		close(resume)
+
+		require.ErrorIs(t, <-queueErr, ErrClientShutdown,
+			"a message accepted after the drain would never be published or called back")
+		assert.False(t, callbackRan.Load())
+	})
 }
 
 func countCounters(counters *sync.Map) int {