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..b03a5eed5d 100644 --- a/pkg/chipingress/batch/client.go +++ b/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,10 +209,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 @@ -216,7 +221,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. @@ -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 diff --git a/pkg/chipingress/batch/client_test.go b/pkg/chipingress/batch/client_test.go index 31c0aa5ca9..a7c269687e 100644 --- a/pkg/chipingress/batch/client_test.go +++ b/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 {