From bb9db88a736e6c016229075e55439da060ea4ca9 Mon Sep 17 00:00:00 2001 From: Chet Sharma Date: Tue, 13 Oct 2026 13:21:33 -0700 Subject: [PATCH 1/3] chipingress/batch: publish queued messages on Stop instead of discarding them On Stop the batcher flushed only the batch it was filling. Messages still queued on the buffer channel, typically when sends are saturated, were abandoned: never published and their callbacks never ran, so callers could not count them as dropped or failed. The batcher now drains everything already queued, in batches of at most batchSize, before returning. Stop's shutdown timeout still bounds the wait; batches drained after it expires fail against the closed client and report the error to their callbacks. --- pkg/chipingress/batch/batcher.go | 21 +++++++- pkg/chipingress/batch/client.go | 5 +- pkg/chipingress/batch/client_test.go | 76 ++++++++++++++++++++++++++++ 3 files changed, 98 insertions(+), 4 deletions(-) diff --git a/pkg/chipingress/batch/batcher.go b/pkg/chipingress/batch/batcher.go index c2c50c9308..92ac457ca2 100644 --- a/pkg/chipingress/batch/batcher.go +++ b/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) + } +} diff --git a/pkg/chipingress/batch/client.go b/pkg/chipingress/batch/client.go index e912b21f8a..54868debb5 100644 --- a/pkg/chipingress/batch/client.go +++ b/pkg/chipingress/batch/client.go @@ -205,10 +205,11 @@ 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 diff --git a/pkg/chipingress/batch/client_test.go b/pkg/chipingress/batch/client_test.go index 31c0aa5ca9..774ee182ed 100644 --- a/pkg/chipingress/batch/client_test.go +++ b/pkg/chipingress/batch/client_test.go @@ -18,6 +18,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 +1216,81 @@ 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) + }) } func countCounters(counters *sync.Map) int { From d8a7031d59ad1494b24a5042184848d40eabb73f Mon Sep 17 00:00:00 2001 From: Chet Sharma Date: Tue, 13 Oct 2026 13:46:25 -0700 Subject: [PATCH 2/3] chipingress/batch: close enqueue race with Stop; test the timeout path QueueMessage checked stopCh and enqueued later, so a producer could add a message after the batcher had drained and returned; that message was never published and its callback never ran. An admission RWMutex now makes the check-and-enqueue atomic with Stop closing stopCh. Add tests for the shutdown-timeout path (Stop returns at the bound, closes the client, late batches report errors to every callback) and for producers racing Stop (every accepted message gets a callback). --- pkg/chipingress/batch/client.go | 30 ++++--- pkg/chipingress/batch/client_test.go | 115 +++++++++++++++++++++++++++ 2 files changed, 136 insertions(+), 9 deletions(-) diff --git a/pkg/chipingress/batch/client.go b/pkg/chipingress/batch/client.go index 54868debb5..9a423dd990 100644 --- a/pkg/chipingress/batch/client.go +++ b/pkg/chipingress/batch/client.go @@ -103,14 +103,17 @@ 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 + 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 @@ -217,7 +220,9 @@ func (b *Client) Stop() { 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. @@ -280,7 +285,7 @@ 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 @@ -313,6 +318,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 diff --git a/pkg/chipingress/batch/client_test.go b/pkg/chipingress/batch/client_test.go index 774ee182ed..8da722929c 100644 --- a/pkg/chipingress/batch/client_test.go +++ b/pkg/chipingress/batch/client_test.go @@ -3,10 +3,12 @@ package batch import ( "context" "errors" + "fmt" "slices" "strconv" "strings" "sync" + "sync/atomic" "testing" "time" @@ -1291,6 +1293,119 @@ func TestStop(t *testing.T) { 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("every message accepted while Stop races producers gets a callback", func(t *testing.T) { + mockClient := mocks.NewClient(t) + mockClient.EXPECT().Close().Return(nil).Maybe() + mockClient.EXPECT().PublishBatch(mock.Anything, mock.Anything). + RunAndReturn(func(_ context.Context, batch *chipingress.CloudEventBatch, _ ...grpc.CallOption) (*chipingress.PublishResponse, error) { + 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 + }).Maybe() + + client, err := NewBatchClient(mockClient, WithBatchSize(50), WithMessageBuffer(100_000), WithBatchInterval(time.Millisecond)) + require.NoError(t, err) + client.Start(t.Context()) + + var accepted, callbacks atomic.Int64 + var producers sync.WaitGroup + for p := range 8 { + producers.Go(func() { + for i := 0; ; i++ { + err := client.QueueMessage(&chipingress.CloudEventPb{ + Id: fmt.Sprintf("%d-%d", p, i), Source: "test-source", Type: "test.event.type", + }, func(error) { callbacks.Add(1) }) + if errors.Is(err, ErrClientShutdown) { + return + } + if err == nil { + accepted.Add(1) + } + } + }) + } + time.Sleep(20 * time.Millisecond) + client.Stop() + producers.Wait() + + require.Positive(t, accepted.Load()) + assert.Equal(t, accepted.Load(), callbacks.Load(), "a message QueueMessage accepted was abandoned") + }) } func countCounters(counters *sync.Map) int { From 456da1f7ac1b163df1d1e6797b959e184ca34cae Mon Sep 17 00:00:00 2001 From: Chet Sharma Date: Tue, 13 Oct 2026 14:32:59 -0700 Subject: [PATCH 3/3] chipingress/batch: test the Stop admission race deterministically The producers-vs-Stop stress test passed with the admission lock removed: the window between the fast-path stop check and the enqueue is too narrow to hit by chance. Replace it with a test that parks a producer in that window through a per-client test hook (nil in production), runs Stop, then resumes the producer and requires ErrClientShutdown. It fails without the admission re-check. --- pkg/chipingress/batch/client.go | 4 +++ pkg/chipingress/batch/client_test.go | 54 ++++++++++++---------------- 2 files changed, 26 insertions(+), 32 deletions(-) diff --git a/pkg/chipingress/batch/client.go b/pkg/chipingress/batch/client.go index 9a423dd990..b03a5eed5d 100644 --- a/pkg/chipingress/batch/client.go +++ b/pkg/chipingress/batch/client.go @@ -106,6 +106,7 @@ type Client struct { // 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 @@ -291,6 +292,9 @@ func (b *Client) QueueMessage(event *chipingress.CloudEventPb, callback func(err return ErrClientShutdown default: } + if b.afterFastPath != nil { + b.afterFastPath() + } eventToQueue := event if b.cloneEvent { diff --git a/pkg/chipingress/batch/client_test.go b/pkg/chipingress/batch/client_test.go index 8da722929c..a7c269687e 100644 --- a/pkg/chipingress/batch/client_test.go +++ b/pkg/chipingress/batch/client_test.go @@ -3,7 +3,6 @@ package batch import ( "context" "errors" - "fmt" "slices" "strconv" "strings" @@ -1366,45 +1365,36 @@ func TestStop(t *testing.T) { assert.Equal(t, queued-2, failed, "batches drained after the timeout report the closed-client error") }) - t.Run("every message accepted while Stop races producers gets a callback", func(t *testing.T) { + 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() - mockClient.EXPECT().PublishBatch(mock.Anything, mock.Anything). - RunAndReturn(func(_ context.Context, batch *chipingress.CloudEventBatch, _ ...grpc.CallOption) (*chipingress.PublishResponse, error) { - 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 - }).Maybe() - client, err := NewBatchClient(mockClient, WithBatchSize(50), WithMessageBuffer(100_000), WithBatchInterval(time.Millisecond)) + client, err := NewBatchClient(mockClient, WithBatchInterval(time.Hour)) require.NoError(t, err) client.Start(t.Context()) - var accepted, callbacks atomic.Int64 - var producers sync.WaitGroup - for p := range 8 { - producers.Go(func() { - for i := 0; ; i++ { - err := client.QueueMessage(&chipingress.CloudEventPb{ - Id: fmt.Sprintf("%d-%d", p, i), Source: "test-source", Type: "test.event.type", - }, func(error) { callbacks.Add(1) }) - if errors.Is(err, ErrClientShutdown) { - return - } - if err == nil { - accepted.Add(1) - } - } - }) + reachedHook := make(chan struct{}) + resume := make(chan struct{}) + client.afterFastPath = func() { + close(reachedHook) + <-resume // hold the producer between the fast-path check and admission } - time.Sleep(20 * time.Millisecond) - client.Stop() - producers.Wait() - require.Positive(t, accepted.Load()) - assert.Equal(t, accepted.Load(), callbacks.Load(), "a message QueueMessage accepted was abandoned") + 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()) }) }