Update state when batch dropped (#1789)

duhminick · web-flow · commit 2da9c430a4bb · 2025-08-13T02:02:31.000-05:00
diff --git a/plugins/outputs/cloudwatchlogs/internal/pusher/batch.go b/plugins/outputs/cloudwatchlogs/internal/pusher/batch.go
@@ -98,7 +98,9 @@ type logEventBatch struct {
 	minT, maxT time.Time
 	// Callbacks to execute when batch is successfully sent.
 	doneCallbacks []func()
-	batchers      map[string]*state.RangeQueueBatcher
+	// Callbacks specifically for updating state
+	stateCallbacks []func()
+	batchers       map[string]*state.RangeQueueBatcher
 }
 
 func newLogEventBatch(target Target, entityProvider logs.LogEntityProvider) *logEventBatch {
@@ -151,7 +153,7 @@ func (b *logEventBatch) handleLogEventState(s *logEventState) {
 	batcher, ok := b.batchers[queueID]
 	if !ok {
 		batcher = state.NewRangeQueueBatcher(s.queue)
-		b.addDoneCallback(batcher.Done)
+		b.addStateCallback(batcher.Done)
 		b.batchers[queueID] = batcher
 	}
 	batcher.Merge(s.r)
@@ -164,14 +166,36 @@ func (b *logEventBatch) addDoneCallback(callback func()) {
 	}
 }
 
-// done runs all registered callbacks.
+// addStateCallback adds the callback to the state callbacks list.
+// State callbacks are specifically for updating the state file and are executed
+// even when a batch fails after exhausting all retry attempts.
+func (b *logEventBatch) addStateCallback(callback func()) {
+	if callback != nil {
+		b.stateCallbacks = append(b.stateCallbacks, callback)
+	}
+}
+
+// done runs all registered callbacks, including both success callbacks and state callbacks.
 func (b *logEventBatch) done() {
+	b.updateState()
+
 	for i := len(b.doneCallbacks) - 1; i >= 0; i-- {
 		done := b.doneCallbacks[i]
 		done()
 	}
 }
 
+// updateState runs only the state callbacks to update the state file
+// without executing other success-related callbacks. This is used when a batch
+// fails after exhausting all retry attempts to prevent reprocessing the same
+// batch after restart.
+func (b *logEventBatch) updateState() {
+	for i := len(b.stateCallbacks) - 1; i >= 0; i-- {
+		callback := b.stateCallbacks[i]
+		callback()
+	}
+}
+
 // build creates a cloudwatchlogs.PutLogEventsInput from the batch. The log events in the batch must be in
 // chronological order by their timestamp.
 func (b *logEventBatch) build() *cloudwatchlogs.PutLogEventsInput {
diff --git a/plugins/outputs/cloudwatchlogs/internal/pusher/batch_test.go b/plugins/outputs/cloudwatchlogs/internal/pusher/batch_test.go
@@ -51,6 +51,137 @@ func TestLogEvent(t *testing.T) {
 }
 
 func TestLogEventBatch(t *testing.T) {
+	t.Run("UpdateStateOnly", func(t *testing.T) {
+		batch := newLogEventBatch(Target{Group: "G", Stream: "S"}, nil)
+
+		successCallbackCalled := false
+		successCallback := func() {
+			successCallbackCalled = true
+		}
+
+		stateCallbackCalled := false
+		stateCallback := func() {
+			stateCallbackCalled = true
+		}
+
+		batch.addDoneCallback(successCallback)
+		batch.addStateCallback(stateCallback)
+
+		batch.updateState()
+
+		assert.False(t, successCallbackCalled, "Success callback should not have been called")
+		assert.True(t, stateCallbackCalled, "State callback should have been called")
+	})
+
+	t.Run("UpdateStateOnly_WithMultipleCallbacks", func(t *testing.T) {
+		batch := newLogEventBatch(Target{Group: "G", Stream: "S"}, nil)
+
+		successCallbacksCalled := make([]bool, 3)
+		successCallbacks := []func(){
+			func() { successCallbacksCalled[0] = true },
+			func() { successCallbacksCalled[1] = true },
+			func() { successCallbacksCalled[2] = true },
+		}
+
+		stateCallbacksCalled := make([]bool, 3)
+		stateCallbacks := []func(){
+			func() { stateCallbacksCalled[0] = true },
+			func() { stateCallbacksCalled[1] = true },
+			func() { stateCallbacksCalled[2] = true },
+		}
+
+		for _, cb := range successCallbacks {
+			batch.addDoneCallback(cb)
+		}
+		for _, cb := range stateCallbacks {
+			batch.addStateCallback(cb)
+		}
+
+		batch.updateState()
+
+		// Verify none of the success callbacks were called
+		for i, called := range successCallbacksCalled {
+			assert.False(t, called, "Success callback %d should not have been called", i)
+		}
+
+		// Verify all state callbacks were called
+		for i, called := range stateCallbacksCalled {
+			assert.True(t, called, "State callback %d should have been called", i)
+		}
+	})
+
+	t.Run("UpdateStateOnly_WithRangeQueueBatcher", func(t *testing.T) {
+		batch := newLogEventBatch(Target{Group: "G", Stream: "S"}, nil)
+
+		mrq1 := &mockRangeQueue{}
+		mrq1.On("ID").Return("test1")
+		mrq1.On("Enqueue", state.NewRange(10, 20)).Once()
+
+		mrq2 := &mockRangeQueue{}
+		mrq2.On("ID").Return("test2")
+		mrq2.On("Enqueue", state.NewRange(30, 40)).Once()
+
+		event1 := newStatefulLogEvent(time.Now(), "Test1", nil, &logEventState{
+			r:     state.NewRange(10, 20),
+			queue: mrq1,
+		})
+		event2 := newStatefulLogEvent(time.Now(), "Test2", nil, &logEventState{
+			r:     state.NewRange(30, 40),
+			queue: mrq2,
+		})
+
+		successCallbackCalled := false
+		batch.addDoneCallback(func() {
+			successCallbackCalled = true
+		})
+
+		batch.append(event1)
+		batch.append(event2)
+
+		batch.updateState()
+
+		mrq1.AssertExpectations(t)
+		mrq2.AssertExpectations(t)
+
+		assert.False(t, successCallbackCalled, "Success callback should not have been called")
+	})
+
+	t.Run("UpdateStateWithDone", func(t *testing.T) {
+		batch := newLogEventBatch(Target{Group: "G", Stream: "S"}, nil)
+
+		mrq1 := &mockRangeQueue{}
+		mrq1.On("ID").Return("test1")
+		mrq1.On("Enqueue", state.NewRange(10, 20)).Once()
+
+		mrq2 := &mockRangeQueue{}
+		mrq2.On("ID").Return("test2")
+		mrq2.On("Enqueue", state.NewRange(30, 40)).Once()
+
+		event1 := newStatefulLogEvent(time.Now(), "Test1", nil, &logEventState{
+			r:     state.NewRange(10, 20),
+			queue: mrq1,
+		})
+		event2 := newStatefulLogEvent(time.Now(), "Test2", nil, &logEventState{
+			r:     state.NewRange(30, 40),
+			queue: mrq2,
+		})
+
+		stateCallbackCalled := false
+		batch.addStateCallback(func() {
+			stateCallbackCalled = true
+		})
+
+		batch.append(event1)
+		batch.append(event2)
+
+		batch.done()
+
+		mrq1.AssertExpectations(t)
+		mrq2.AssertExpectations(t)
+
+		assert.True(t, stateCallbackCalled, "State callback should have been called")
+	})
+
 	t.Run("Append", func(t *testing.T) {
 		batch := newLogEventBatch(Target{Group: "G", Stream: "S"}, nil)
 
diff --git a/plugins/outputs/cloudwatchlogs/internal/pusher/queue_test.go b/plugins/outputs/cloudwatchlogs/internal/pusher/queue_test.go
@@ -15,8 +15,11 @@ import (
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/aws/awserr"
 	"github.com/influxdata/telegraf"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
 
+	"github.com/aws/amazon-cloudwatch-agent/internal/state"
 	"github.com/aws/amazon-cloudwatch-agent/logs"
 	"github.com/aws/amazon-cloudwatch-agent/sdk/service/cloudwatchlogs"
 	"github.com/aws/amazon-cloudwatch-agent/tool/testutil"
@@ -66,6 +69,23 @@ func (s *stubLogsService) DescribeLogGroups(in *cloudwatchlogs.DescribeLogGroups
 	return nil, nil
 }
 
+type mockSender struct {
+	mock.Mock
+}
+
+func (m *mockSender) Send(batch *logEventBatch) {
+	m.Called(batch)
+}
+
+func (m *mockSender) SetRetryDuration(d time.Duration) {
+	m.Called(d)
+}
+
+func (m *mockSender) RetryDuration() time.Duration {
+	args := m.Called()
+	return args.Get(0).(time.Duration)
+}
+
 func TestAddSingleEvent_WithAccountId(t *testing.T) {
 	t.Parallel()
 	var wg sync.WaitGroup
@@ -685,3 +705,98 @@ func testPreparationWithLogger(
 	)
 	return stop, q.(*queue)
 }
+
+func TestQueueCallbackRegistration(t *testing.T) {
+	t.Run("RegistersCallbacks", func(t *testing.T) {
+		var wg sync.WaitGroup
+		var s stubLogsService
+		var called bool
+
+		// Mock the PutLogEvents method to verify the batch has callbacks registered
+		s.ple = func(*cloudwatchlogs.PutLogEventsInput) (*cloudwatchlogs.PutLogEventsOutput, error) {
+			called = true
+			return &cloudwatchlogs.PutLogEventsOutput{}, nil
+		}
+
+		mockSender := &mockSender{}
+		mockSender.On("Send", mock.AnythingOfType("*pusher.logEventBatch")).Run(func(args mock.Arguments) {
+			batch := args.Get(0).(*logEventBatch)
+
+			assert.NotEmpty(t, batch.doneCallbacks, "Regular callbacks should be registered")
+			assert.Empty(t, batch.stateCallbacks, "State callbacks should not be registered")
+
+			s.PutLogEvents(batch.build())
+		}).Return()
+
+		logger := testutil.NewNopLogger()
+		stop := make(chan struct{})
+		q := &queue{
+			target:          Target{"G", "S", util.StandardLogGroupClass, -1},
+			logger:          logger,
+			converter:       newConverter(logger, Target{"G", "S", util.StandardLogGroupClass, -1}),
+			batch:           newLogEventBatch(Target{"G", "S", util.StandardLogGroupClass, -1}, nil),
+			sender:          mockSender,
+			eventsCh:        make(chan logs.LogEvent, 100),
+			flushCh:         make(chan struct{}),
+			resetTimerCh:    make(chan struct{}),
+			flushTimer:      time.NewTimer(10 * time.Millisecond),
+			stop:            stop,
+			startNonBlockCh: make(chan struct{}),
+			wg:              &wg,
+		}
+		q.flushTimeout.Store(10 * time.Millisecond)
+
+		q.batch.append(newLogEvent(time.Now(), "test message", nil))
+		q.send()
+
+		mockSender.AssertExpectations(t)
+		assert.True(t, called, "PutLogEvents should have been called")
+	})
+
+	t.Run("RegistersStateCallbacksForStatefulEvents", func(t *testing.T) {
+		var wg sync.WaitGroup
+
+		mrq := &mockRangeQueue{}
+		mrq.On("ID").Return("test-queue")
+		mrq.On("Enqueue", mock.Anything).Return()
+
+		mockSender := &mockSender{}
+		mockSender.On("Send", mock.AnythingOfType("*pusher.logEventBatch")).Run(func(args mock.Arguments) {
+			batch := args.Get(0).(*logEventBatch)
+
+			assert.NotEmpty(t, batch.doneCallbacks, "Regular callbacks should be registered")
+			assert.NotEmpty(t, batch.stateCallbacks, "State callbacks should be registered")
+
+			batcher, ok := batch.batchers["test-queue"]
+			assert.True(t, ok, "Batch should have a batcher for our queue")
+			assert.NotNil(t, batcher, "Batcher should not be nil")
+		}).Return()
+
+		logger := testutil.NewNopLogger()
+		stop := make(chan struct{})
+		q := &queue{
+			target:          Target{"G", "S", util.StandardLogGroupClass, -1},
+			logger:          logger,
+			converter:       newConverter(logger, Target{"G", "S", util.StandardLogGroupClass, -1}),
+			batch:           newLogEventBatch(Target{"G", "S", util.StandardLogGroupClass, -1}, nil),
+			sender:          mockSender,
+			eventsCh:        make(chan logs.LogEvent, 100),
+			flushCh:         make(chan struct{}),
+			resetTimerCh:    make(chan struct{}),
+			flushTimer:      time.NewTimer(10 * time.Millisecond),
+			stop:            stop,
+			startNonBlockCh: make(chan struct{}),
+			wg:              &wg,
+		}
+		q.flushTimeout.Store(10 * time.Millisecond)
+
+		event := newStubStatefulLogEvent("test message", time.Now(), state.NewRange(10, 20), mrq)
+
+		convertedEvent := q.converter.convert(event)
+		q.batch.append(convertedEvent)
+
+		q.send()
+
+		mockSender.AssertExpectations(t)
+	})
+}
diff --git a/plugins/outputs/cloudwatchlogs/internal/pusher/sender.go b/plugins/outputs/cloudwatchlogs/internal/pusher/sender.go
@@ -87,6 +87,7 @@ func (s *sender) Send(batch *logEventBatch) {
 		var awsErr awserr.Error
 		if !errors.As(err, &awsErr) {
 			s.logger.Errorf("Non aws error received when sending logs to %v/%v: %v. CloudWatch agent will not retry and logs will be missing!", batch.Group, batch.Stream, err)
+			batch.updateState()
 			return
 		}
 
@@ -99,6 +100,7 @@ func (s *sender) Send(batch *logEventBatch) {
 		case *cloudwatchlogs.InvalidParameterException,
 			*cloudwatchlogs.DataAlreadyAcceptedException:
 			s.logger.Errorf("%v, will not retry the request", e)
+			batch.updateState()
 			return
 		default:
 			s.logger.Errorf("Aws error received when sending logs to %v/%v: %v", batch.Group, batch.Stream, awsErr)
@@ -116,6 +118,7 @@ func (s *sender) Send(batch *logEventBatch) {
 
 		if time.Since(startTime)+wait > s.RetryDuration() {
 			s.logger.Errorf("All %v retries to %v/%v failed for PutLogEvents, request dropped.", retryCountShort+retryCountLong-1, batch.Group, batch.Stream)
+			batch.updateState()
 			return
 		}
 
@@ -124,6 +127,7 @@ func (s *sender) Send(batch *logEventBatch) {
 		select {
 		case <-s.stop:
 			s.logger.Errorf("Stop requested after %v retries to %v/%v failed for PutLogEvents, request dropped.", retryCountShort+retryCountLong-1, batch.Group, batch.Stream)
+			batch.updateState()
 			return
 		case <-time.After(wait):
 		}
diff --git a/plugins/outputs/cloudwatchlogs/internal/pusher/sender_test.go b/plugins/outputs/cloudwatchlogs/internal/pusher/sender_test.go
